import torch.nn as nn
import torch
class lbn(nn.Module):
    def __init__(self,batchsize):
        super(lbn,self).__init__()
        self.batchsize = batchsize
        self.tensorcache=[]
        self.p=nn.Parameter(torch.tensor(0.5).float())
    def forward(self,x):
        if len(self.tensorcache)>1:
            if x.shape == self.tensorcache[0].shape:
                self.tensorcache.append(x.clone().detach())
        if len(self.tensorcache) < 3:
            return x
        if len(self.tensorcache)>self.batchsize:
            self.tensorcache.pop(0)
        stacked_tensor = torch.stack(self.tensorcache)
        mean = torch.mean(stacked_tensor, dim=0)
        std = torch.std(stacked_tensor, dim=0)
        return (x-mean)/std